#include "kernel_operator.h"
#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;
class KernelDiv
{
public:
  __aicore__ inline KernelDiv() {}
  __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, uint32_t blockLength,
                              uint32_t tileNum, uint32_t tileLength,
                              uint32_t lasttileLength, uint32_t typeKey)
  {
    ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
    this->typeKey = typeKey;
    this->blockLength = blockLength;
    this->tileNum =
        tileNum ASSERT(tileNum != 0 && "tile num can not be zero!");
    this->tileLength = tileLength / BUFFER_NUM;
    this->lasttileLength = lasttileLength;
    x1Gm.SetGlobalBuffer((__gm__ DTYPE_X1 *)x1,
                        this->blockLength);
    x2Gm.SetGlobalBuffer((__gm__ DTYPE_X1 *)x2,
                        this->blockLength);
    yGm.SetGlobalBuffer((__gm__ DTYPE_X1 *)y,
                          this->blockLength);
    pipe.InitBuffer(inQueueX1, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X1));
    pipe.InitBuffer(inQueueX2, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X1));
    pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X1));
    
    if(typeKey == 0 || typeKey == 3){// float int32
      pipe.InitBuffer(calcBuf, 2 * this->tileLength * sizeof(float));
      pipe.InitBuffer(calcBuf2, this->tileLength * sizeof(float));
    }else{
      pipe.InitBuffer(calcBuf, 2 * this->tileLength * sizeof(half));
      pipe.InitBuffer(calcBuf2, this->tileLength * sizeof(half));
    }
  }
  __aicore__ inline void Process()
  {
    int32_t loopCount = this->tileNum * BUFFER_NUM;
    for (int32_t i = 0; i < loopCount; i++)
    {
      CopyIn(i);
      Compute(i);
      CopyOut(i);
    }
  }

private:
  __aicore__ inline void CopyIn(int32_t progress)
  {
    LocalTensor<DTYPE_X1> x1Local = inQueueX1.AllocTensor<DTYPE_X1>();
    LocalTensor<DTYPE_X1> x2Local = inQueueX2.AllocTensor<DTYPE_X1>();

    if (BUFFER_NUM == 1) {
      if (progress == this->tileNum - 1) {
        if (progress == 0) {
          //如果只有一包，则搬运的起始地址为0，tileLength为实际分块的数据量
          DataCopy(x1Local[0], x1Gm[0], this->tileLength);
          DataCopy(x2Local[0], x2Gm[0], this->tileLength);
        } else {
          //将最后一个分块的起始地址向前移动tileLength-lasttileLength
          DataCopy(
              x1Local[0],
              x1Gm[(progress - 1) * this->tileLength + this->lasttileLength],
              this->tileLength);
          DataCopy(
              x2Local[0],
              x2Gm[(progress - 1) * this->tileLength + this->lasttileLength],
              this->tileLength);
        }
      } else {
        DataCopy(x1Local[0], x1Gm[progress * this->tileLength],
                 this->tileLength);
        DataCopy(x2Local[0], x2Gm[progress * this->tileLength],
                 this->tileLength);
        
      }
    }
    if (BUFFER_NUM == 2) {
      if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
          (progress == (this->tileNum * BUFFER_NUM - 1))) {
        DataCopy(
            x1Local[0],
            x1Gm[(progress - 2) * (this->tileLength) + this->lasttileLength],
            (this->tileLength));
        DataCopy(
            x2Local[0],
            x2Gm[(progress - 2) * (this->tileLength) + this->lasttileLength],
            (this->tileLength));

      }
      else {
        DataCopy(x1Local[0], x1Gm[progress * (this->tileLength)],
                 (this->tileLength));
        DataCopy(x2Local[0], x2Gm[progress * (this->tileLength)],
                 (this->tileLength));
      }
    }
    inQueueX1.EnQue(x1Local);
    inQueueX2.EnQue(x2Local);
    
  }
  __aicore__ inline void Compute(int32_t progress)
  {
    if(this->typeKey == 2){ // int8
      LocalTensor<int8_t> x1Local = inQueueX1.DeQue<int8_t>();
      LocalTensor<int8_t> x2Local = inQueueX2.DeQue<int8_t>();
      LocalTensor<int8_t> yLocal = outQueueY.AllocTensor<int8_t>();
      LocalTensor<half> c2Local = calcBuf2.Get<half>();
      LocalTensor<half> tmpLocal = calcBuf.Get<half>();
      // compute
      Cast(tmpLocal, x1Local, RoundMode::CAST_NONE, this->tileLength);
      Cast(tmpLocal[this->tileLength], x2Local, RoundMode::CAST_NONE, this->tileLength);
      Div(c2Local, tmpLocal, tmpLocal[this->tileLength], this->tileLength);
      Cast(yLocal, c2Local, RoundMode::CAST_FLOOR, this->tileLength);
      outQueueY.EnQue<int8_t>(yLocal);
      inQueueX1.FreeTensor(x1Local);
      inQueueX2.FreeTensor(x2Local);
    } else if(this->typeKey == 3){ // int32
      LocalTensor<int32_t> x1Local = inQueueX1.DeQue<int32_t>();
      LocalTensor<int32_t> x2Local = inQueueX2.DeQue<int32_t>();
      LocalTensor<float> c2Local = calcBuf2.Get<float>();
      LocalTensor<int32_t> yLocal = outQueueY.AllocTensor<int32_t>();
      LocalTensor<float> tmpLocal = calcBuf.Get<float>();
      // compute
      Cast(tmpLocal, x1Local, RoundMode::CAST_NONE, this->tileLength);
      Cast(tmpLocal[this->tileLength], x2Local, RoundMode::CAST_NONE, this->tileLength);
      Div(c2Local, tmpLocal, tmpLocal[this->tileLength], this->tileLength);
      Cast(yLocal, c2Local, RoundMode::CAST_FLOOR, this->tileLength);
      outQueueY.EnQue<int32_t>(yLocal);
      inQueueX1.FreeTensor(x1Local);
      inQueueX2.FreeTensor(x2Local);
    }
    else if(this->typeKey == 1){// half
      LocalTensor<half> x1Local = inQueueX1.DeQue<half>();
      LocalTensor<half> x2Local = inQueueX2.DeQue<half>();
      LocalTensor<half> yLocal = outQueueY.AllocTensor<half>();
      // compute
      Div(yLocal, x1Local, x2Local, this->tileLength);
      outQueueY.EnQue<half>(yLocal);
      inQueueX1.FreeTensor(x1Local);
      inQueueX2.FreeTensor(x2Local);
    }
    else{
      LocalTensor<float> x1Local = inQueueX1.DeQue<float>();
      LocalTensor<float> x2Local = inQueueX2.DeQue<float>();
      LocalTensor<float> yLocal = outQueueY.AllocTensor<float>();

      Div(yLocal, x1Local, x2Local, this->tileLength);
      
      outQueueY.EnQue<float>(yLocal);
      inQueueX1.FreeTensor(x1Local);
      inQueueX2.FreeTensor(x2Local);
    }
  }
  __aicore__ inline void CopyOut(int32_t progress)
  {
    LocalTensor<DTYPE_X1> yLocal = outQueueY.DeQue<DTYPE_X1>();

    if (BUFFER_NUM == 1) {
      if (progress == this->tileNum - 1) {
        if (progress == 0) {
          DataCopy(yGm[0], yLocal, this->tileLength);
        } else {
          DataCopy(
              yGm[(progress - 1) * this->tileLength + this->lasttileLength],
              yLocal, this->tileLength);
        }
      } else {
        DataCopy(yGm[progress * this->tileLength], yLocal,
                 this->tileLength);
      }
    }
    if (BUFFER_NUM == 2) {
      if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
          (progress == (this->tileNum * BUFFER_NUM - 1))) {
        DataCopy(
            yGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
            yLocal, (this->tileLength));
      }
      else {
        DataCopy(yGm[progress * (this->tileLength)], yLocal,
                 (this->tileLength));
      }
    }

    outQueueY.FreeTensor(yLocal);
  }

private:
  TPipe pipe;
  TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX1, inQueueX2;
  TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
  TBuf<TPosition::VECCALC> calcBuf, calcBuf2;
  GlobalTensor<DTYPE_X1> x1Gm;
  GlobalTensor<DTYPE_X1> x2Gm;
  GlobalTensor<DTYPE_X1> yGm;
  uint32_t blockLength;
  uint32_t tileNum;
  uint32_t tileLength;
  uint32_t lasttileLength;
  uint32_t typeKey;
};


extern "C" __global__ __aicore__ void div(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    KernelDiv op;

    op.Init(x1, x2, y, tiling_data.blockLength,
            tiling_data.tileNum, tiling_data.tileLength,
            tiling_data.lasttileLength, tiling_data.typeKey);
    op.Process();
}